import codecs

special_symbol = [' <A0>', '<A0> ', ' <V>', ' <A1>', ' <A2>', ' <S>']

prompt_path = '/home/rickwwang/project_research/fairseq/data-bin/prompt.txt.test'
all_prompts = codecs.open(prompt_path, 'r', encoding='utf8').readlines()
all_prompts = map(lambda x: x.strip(), all_prompts)


def is_fit(story):
    return story

exp_file = []
exp_file.append('out/bilstm_transformer-lm_bpe500_format_pre_srl_hier_pos_fix/17hyp.txt.own.prompt2srl2_500_wo_smooth_d1_verb_loss.1.topk10.t1.0')


base_path = '/home/rickwwang/project_research/'

story_path = base_path + exp_file[0]
result_path = story_path + '.human.test.format'
all_stories = codecs.open(story_path, 'r', encoding='utf8').readlines()

items = []
for s in all_stories:
    s = s.strip()
    key = s.split(' <SEP> ')[0]
    value = s.split(' <SEP> ')[1]

    verbs = value.split(' <V> ')[1:]
    verbs = [v.split(' ')[0] for v in verbs]
    for sp in special_symbol:
        key = key.replace(sp, '')
    event_words = key.split(' ')
    story_words = value.split(' ')

    verb_len = len(verbs) if len(verbs) > 0 else 1
    event_len = len(event_words)
    story_len = len(story_words)

    event_unique_verb = set(verbs)
    event_unique_1gram = set(event_words)
    event_unique_2gram = set([tuple(event_words[i:i + 2]) for i in range(event_len - 1)])
    story_unique_1gram = set(story_words)
    story_unique_2gram = set([tuple(story_words[i:i + 2]) for i in range(story_len - 1)])
    event_unique_verb = len(event_unique_verb) * 1.0 / verb_len
    event_unique_1gram = len(event_unique_1gram) * 1.0 / event_len
    event_unique_2gram = len(event_unique_2gram) * 1.0 / event_len
    story_unique_1gram = len(story_unique_1gram) * 1.0 / story_len
    story_unique_2gram = len(story_unique_2gram) * 1.0 / story_len

    item = (key, value, event_unique_verb, event_unique_1gram, event_unique_2gram, story_unique_1gram, story_unique_2gram)
    items.append(item)

items = sorted(items, key=lambda x: (x[2], x[6], x[5], x[4], x[3]), reverse=True)
items = items[:1000]

promp_to_story = dict()
for s in items:
    key = s[0]
    value = s[1]

    if not is_fit(s[2:]):
        continue
    if key in promp_to_story:
        promp_to_story[key].append(value)
    else:
        promp_to_story[key] = [value]

format_stories = []
for p in all_prompts:
    if p not in promp_to_story:
        continue
    ss = promp_to_story[p]

    format_stories.append(p)
    format_stories.append(ss[0])

with codecs.open(result_path, 'w', encoding='utf8') as fout:
    fout.write('\n'.join(format_stories))
    fout.write('\n')